"""Benchmark for BINARY_OP_EXTEND type propagation.

Tests whether the tier 2 optimizer can eliminate guards when types are
known from previous BINARY_OP_EXTEND results.

Usage:
    ./python bench_binary_op_extend.py
    ./python bench_binary_op_extend.py --save result.json
    ./python bench_binary_op_extend.py --compare a.json b.json
"""

import sys
import pyperf

INNER = 2000


def bench_list_concat_subscr(n):
    """list + list followed by subscript — tests list type propagation."""
    a = [1, 2, 3]
    b = [4, 5, 6]
    total = 0
    for _ in range(n):
        c = a + b
        total += c[0] + c[3]
    return total


def bench_tuple_concat_unpack(n):
    """tuple + tuple followed by unpack — tests tuple type propagation."""
    t1 = (1, 2)
    t2 = (3, 4)
    total = 0
    for _ in range(n):
        a, b, c, d = t1 + t2
        total += a + d
    return total


def bench_str_repeat(n):
    """str * int in a loop — tests str type propagation."""
    s = "ab"
    total = 0
    for i in range(n):
        r = s * (i % 5)
        total += len(r)
    return total


def bench_bytes_concat(n):
    """bytes + bytes in a loop — tests bytes type propagation."""
    a = b"hello"
    b_ = b" world"
    total = 0
    for _ in range(n):
        c = a + b_
        total += len(c)
    return total


def bench_bytes_repeat(n):
    """bytes * int in a loop — tests bytes type propagation."""
    b = b"ab"
    total = 0
    for i in range(n):
        r = b * (i % 3)
        total += len(r)
    return total


def bench_tuple_repeat(n):
    """tuple * int in a loop — tests tuple type propagation."""
    t = (1, 2, 3)
    total = 0
    for i in range(n):
        r = t * (i % 3)
        total += len(r)
    return total


def bench_dict_merge(n):
    """dict | dict in a loop — tests dict type propagation."""
    d1 = {"a": 1, "b": 2}
    d2 = {"c": 3, "d": 4}
    total = 0
    for _ in range(n):
        d = d1 | d2
        total += len(d)
    return total


def bench_chained_list_ops(n):
    """Multiple list ops chained — tests guard elimination across ops."""
    a = [1, 2]
    b = [3, 4]
    total = 0
    for _ in range(n):
        c = a + b
        d = c + a
        total += d[0] + d[4]
    return total


def bench_mixed_float_int(n):
    """float + int and int + float — existing EXTEND specializations."""
    x = 1.5
    total = 0.0
    for i in range(n):
        a = x + i
        total += a 
    return total

def float_mix_mul(n):
    """float + int then float * float — tests unique flag for inplace mul."""
    x = 1.5
    total = 0.0
    for i in range(n):
        a = (x + i) * 2.0  # result of x+i should be unique -> inplace multiply
        total += a
    return total

BENCHMARKS = [
    ("list_concat_subscr", bench_list_concat_subscr),
    ("tuple_concat_unpack", bench_tuple_concat_unpack),
    ("str_repeat", bench_str_repeat),
    ("bytes_concat", bench_bytes_concat),
    ("bytes_repeat", bench_bytes_repeat),
    ("tuple_repeat", bench_tuple_repeat),
    ("dict_merge", bench_dict_merge),
    ("chained_list_ops", bench_chained_list_ops),
    ("mixed_float_int", bench_mixed_float_int),
    ("float_mix_mul", float_mix_mul),
]
    

def main():
    args = sys.argv[1:]

    if "--compare" in args:
        idx = args.index("--compare")
        file_a = args[idx + 1]
        file_b = args[idx + 2]
        import subprocess
        subprocess.run([sys.executable, "-m", "pyperf", "compare_to",
                       file_a, file_b, "--table"])
        return

    save_file = None
    if "--save" in args:
        idx = args.index("--save")
        save_file = args[idx + 1]

    runner = pyperf.Runner()
    for name, func in BENCHMARKS:
        # Warm up
        func(INNER)
        runner.bench_func(name, func, INNER)

    if save_file and runner.args.output:
        import shutil
        shutil.copy(runner.args.output, save_file)


if __name__ == "__main__":
    main()